package org.custom.spring.boot.datasource.config;

import cn.hutool.core.collection.ConcurrentHashSet;
import cn.hutool.core.lang.Snowflake;
import cn.hutool.core.util.StrUtil;
import org.custom.spring.boot.datasource.annotation.DatasourceSelection;
import org.custom.spring.boot.datasource.application.ApplicationUtil;
import org.custom.spring.boot.datasource.constant.JettExceptionConstant;
import org.custom.spring.boot.datasource.entity.UserInfo;
import org.custom.spring.boot.datasource.entity.datasource.CustomDataSourceProperties;
import org.custom.spring.boot.datasource.entity.datasource.DynamicDatasourceSelection;
import org.custom.spring.boot.datasource.exception.TipException;
import org.custom.spring.boot.datasource.manager.DataSourceManager;
import org.custom.spring.boot.datasource.service.DataSourceBuilder;
import org.custom.spring.boot.datasource.service.DynamicDatasourceService;
import org.custom.spring.boot.datasource.service.UserInfoObtainService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import javax.sql.DataSource;
import javax.validation.Validator;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import static cn.hutool.core.text.CharSequenceUtil.EMPTY;
import static org.custom.spring.boot.datasource.service.DataSourceBuilder.DATA_SOURCE_BUILDERS;

/**
 * 动态数据源配置类
 *
 * @author 李坤
 * @date 2022/3/31 10:08
 */
@Component
public class DynamicRoutingDataSource extends AbstractRoutingDataSource implements DynamicDatasourceService {
    private final Logger logger = LoggerFactory.getLogger(DynamicRoutingDataSource.class);
    /**
     * 动态数据源集合
     */
    private final Map<Object, Object> customDataSources = new ConcurrentHashMap<>();
    /**
     * 数据源KEY与用户信息映射
     */
    private final Map<String, Set<String>> ownerDataSourceMapping = new ConcurrentHashMap<>();
    /**
     * 默认数据源类型
     */
    private Class<? extends DataSource> defaultDataSourceClass;
    @Resource
    private Snowflake snowflake;
    @Resource
    private Validator validator4DataSource;
    @Autowired(required = false)
    private DataSourceProperties dataSourceProperties;

    @PostConstruct
    public void initDynamicRoutingDataSource() {
        // 处理默认数据源
        if (Objects.nonNull(dataSourceProperties)) {
            // 创建并加载默认数据源配置
            DataSource dataSource = dataSourceProperties.initializeDataSourceBuilder().build();

            // 设置默认数据源（单独使用非必须；整合其他项目必须）
            this.setDefaultTargetDataSource(dataSource);
            // 将默认数据源的类型保存下来（便于后续添加数据源）
            this.defaultDataSourceClass = dataSource.getClass();
            // 保存默认数据源，以便于统一处理
            customDataSources.put(DatasourceSelection.DEFAULT, dataSource);
        }
        // 设置动态数据源集合
        this.setTargetDataSources(customDataSources);
    }

    @Override
    protected Object determineCurrentLookupKey() {
        String dataSourceCode = Optional.ofNullable(DataSourceManager.getDynamicDatasourceSelection()).map(DynamicDatasourceSelection::getDataSourceKey).orElse(EMPTY);

        logger.info("当前数据源编码：{}", Objects.nonNull(dataSourceCode) && customDataSources.containsKey(dataSourceCode) ? dataSourceCode : DatasourceSelection.DEFAULT);

        return dataSourceCode;
    }


    @Override
    public Connection obtainConnection(DataSource dataSource) {
        Connection connection;
        try {
            connection = dataSource.getConnection();
        } catch (SQLException exception) {
            connection = null;
        }
        return connection;
    }

    @Override
    public Class<? extends DataSource> getDefaultDataSourceClass() {
        return this.defaultDataSourceClass;
    }

    @Override
    public void setDefaultDataSourceClass(Class<? extends DataSource> dataSourceClass) {
        this.defaultDataSourceClass = dataSourceClass;
    }

    @Override
    public DataSource getDataSource(String dataSourceKey) {
        // 判断数据源是否存在
        if (!customDataSources.containsKey(dataSourceKey)) {
            throw new TipException(String.format(JettExceptionConstant.SOMETHING_NOT_EXIST, "数据源" + dataSourceKey));
        }
        return (DataSource) customDataSources.get(dataSourceKey);
    }

    @Override
    public void verifyDynamicDataSourceSelection(DynamicDatasourceSelection datasourceSelection) {
        // 判断数据源是否存在
        if (!customDataSources.containsKey(datasourceSelection.getDataSourceKey())) {
            if (!datasourceSelection.isUserDefaultDataSource() || Objects.isNull(defaultDataSourceClass)) {
                throw new TipException(String.format(JettExceptionConstant.SOMETHING_NOT_EXIST, "数据源"));
            }
            datasourceSelection.setDataSourceKey(DatasourceSelection.DEFAULT);
        }
        // 判断数据源是否可用
        Optional.ofNullable(this.obtainConnection(this.getDataSource(datasourceSelection.getDataSourceKey())))
                .orElseThrow(() -> new TipException(String.format(JettExceptionConstant.SOMETHING_NOT_EXIST, "数据源")));
    }

    @Override
    public boolean resetDefaultDataSource(CustomDataSourceProperties dataSourceProperties) {
        // 重新设置默认数据类型
        this.defaultDataSourceClass = dataSourceProperties.getType();

        // 根据数据源属性生成数据源
        DataSource dataSource = this.generateDataSource(dataSourceProperties);
        // 设置默认数据源
        super.setDefaultTargetDataSource(dataSource);
        // 此代码是必须的（否则添加的数据源不生效）
        super.afterPropertiesSet();
        return true;
    }

    @Override
    public DataSource generateDataSource(CustomDataSourceProperties dataSourceProperties) {
        // 生成的数据源
        DataSource dataSource;

        if (StrUtil.isNotEmpty(dataSourceProperties.getName())) {
            // 自定义数据源构建服务
            DataSourceBuilder dataSourceBuilder = DATA_SOURCE_BUILDERS.get(defaultDataSourceClass);

            // 判断是否存在自定义的数据源构建服务
            if (Objects.nonNull(dataSourceBuilder)) {
                // 调用自定义数据源构建接口
                dataSource = dataSourceBuilder.generateDataSource(dataSourceProperties);
            } else {
                // 根据数据源参数创建数据源对象
                dataSource = dataSourceProperties.initializeDataSourceBuilder().build();
            }
            return dataSource;
        }
        throw new TipException(String.format(JettExceptionConstant.SOMETHING_NOT_EMPTY, "数据源名称"));
    }

    @Override
    public void appendDataSource(CustomDataSourceProperties dataSourceProperties) {
        // 当前动态数据源ID标识
        String dynamicKey;

        if (StrUtil.isEmpty(dynamicKey = dataSourceProperties.getId())) {
            throw new TipException(String.format(JettExceptionConstant.SOMETHING_NOT_EMPTY, "动态数据源ID"));
        } else if (Objects.equals(DatasourceSelection.DEFAULT, dynamicKey)) {
            throw new TipException(String.format(JettExceptionConstant.SOMETHING_PROHIBITED, "该动态数据源ID"));
        }
        // 存储用户信息与数据源ID映射
        this.appendUserInfoDataSourceMapping(dynamicKey, dataSourceProperties.getOwner());
        // 根据数据源属性生成数据源对象
        DataSource dataSource = this.generateDataSource(dataSourceProperties);
        // 构建数据源对象，并将数据源放置到集合中
        customDataSources.put(dynamicKey, dataSource);
        // 此代码是必须的（否则添加的数据源不生效）
        super.afterPropertiesSet();
    }

    /*-------------------------------------------------private method-------------------------------------------------*/

    /**
     * add userId and dynamicKey to ownerDataSourceMapping
     *
     * @param dynamicKey 数据源ID
     */
    private void appendUserInfoDataSourceMapping(String dynamicKey, String owner) {
        Set<String> dynamicKeys;

        // double-check to get dynamicKeys and put it into ownerDataSourceMapping
        if (Objects.isNull(dynamicKeys = ownerDataSourceMapping.get(owner))) {

            synchronized (owner.intern()) {

                if (Objects.isNull(dynamicKeys = ownerDataSourceMapping.get(owner))) {
                    dynamicKeys = new ConcurrentHashSet<>();
                    ownerDataSourceMapping.put(owner, dynamicKeys);
                }
            }
        }
        dynamicKeys.add(dynamicKey);
    }

    /*--------------------------------------------------static class--------------------------------------------------*/

    public static class DynamicDataSourceUtil {
        /**
         * 将给定的对象转换为字符串表示
         *
         * @param target 给定的对象
         *
         * @return 字符串表示
         */
        public static String toString(String target) {
            return Objects.isNull(target) ? EMPTY : target;
        }

        /**
         * 获取当前登陆用户的ID标识
         *
         * @return 用户的ID标识
         */
        public static String obtainUserId() {
            return DynamicDataSourceUtil.obtainUserInfo().getUserId();
        }

        /**
         * 获取当前登陆用户信息
         *
         * @return 当前登录用户
         */
        public static UserInfo obtainUserInfo() {
            UserInfo userInfo;
            try {
                userInfo = ApplicationUtil.getBean4Class(UserInfoObtainService.class).obtainUserInfo();
            } catch (Exception exception) {
                exception.printStackTrace();
                userInfo = UserInfo.getInstance();
            }
            return userInfo;
        }
    }
}
